import time
from telnetlib1 import Telnet
from datetime import datetime,timedelta

floorNo='10'
apvlan='10'  #根据实际修改为AP使用的Vlan
coreSwIP='172.16.100.254' #改成实际核心交换机的地址
swPwd='Ruijie@123'  #交换机登录密码
encode='ascii'
roomTvFile='rootNo-IPTV.txt' #改成实际存储房号和TV地址的文件名
chkedFile='checkedRooms.txt' #改成实际存储已检索的AP及其Mac的文件名
checkedRooms={}
coreTn=None
floorTn=None
lines=None

#修改为实际的楼层和对应交换机地址
floorSwiches={'10':'172.16.100.109','11':'172.16.100.110','12':'172.16.100.111','13':'172.16.100.105','14':'172.16.100.106','15':'172.16.100.107','16':'172.16.100.101','17':'172.16.100.102','18':'172.16.100.103','87':'172.16.100.113','88':'172.16.100.114','89':'172.16.100.115','84':'172.16.100.119','85':'172.16.100.118','86':'172.16.100.117'}

def connCore():
    global coreSwIP,encode,swPwd
    tn=Telnet(coreSwIP)
    tn.read_until(b'Password:')
    tn.write(f'{swPwd}\nen\n'.encode(encoding=encode))
    tn.read_until(b'Password:')
    tn.write(f'{swPwd}\n'.encode(encode))
    tn.read_until(b'#')
    return tn

def connFloor(floor):
    global encode,swPwd
    tn=Telnet(floorSwiches.get(floor))
    tn.read_until(b'Password:')
    tn.write(f'{swPwd}\nen\n'.encode(encode))
    tn.read_until(b'Password:')
    tn.write(f'{swPwd}\n'.encode(encode))
    tn.read_until(b'#')
    return tn

def getTvMac(ipaddr):
    global coreTn
    coreTn.write(f'show arp {ipaddr}\n'.encode('ascii'))
    data = coreTn.read_until(b'#').decode('ascii').split('\r\n')[2]
    if data.startswith('Internet'):
        return data.split()[3]
    else:
        return False

def getAPMac(tvMac):
    global floorTn,apvlan
    floorTn.write(f'show mac add {tvMac}\n'.encode('ascii'))
    data=floorTn.read_until(b'#').decode('ascii').split('\r\n')[3]
    if data.endswith('#'):
        return False
    else:
        ifNo=data.split()[4]
        floorTn.write(f'show mac interface gigabitEthernet {ifNo}\n'.encode('ascii'))
        data=floorTn.read_until(b'#').decode('ascii').split('\r\n')
        for line in data:
            line=line.split()
            if line[0]==apvlan:
                return line[1]
    return False

def new_func():
    global checkedRooms,floorTn,coreTn,floorNo,lines,floorSwiches,roomTvFile,chkedFile
    coreTn=connCore()
    floorTn=connFloor(floorNo)
    with open(roomTvFile) as reader:
        lines={k:v for k,v in map(lambda x:x.split(),reader.readlines())}
    for k,v in lines.items():  #k是房号，v是电视机IP地址
        tvMac=getTvMac(v)
        if not tvMac: 
            continue
        no=k[0:2]  #获取楼层
        if no!=floorNo:
            floorTn.close()
            floorNo=no
            floorTn=connFloor(floorNo)
        apMac=getAPMac(tvMac)
        if apMac:
            checkedRooms[k]=apMac
        print(f'floor: {k}, ap: {apMac}')
    for room in checkedRooms.keys():
        lines.pop(room)
    with open(roomTvFile,'w') as writer:
        for k,v in lines.items():
            writer.write(f'{k}\t{v}\n')
    with open(chkedFile,'a+') as writer:
        for k,v in checkedRooms.items():
            writer.write(f'{v}\t{k}\n')
    coreTn.close()
    floorTn.close()
    checkedRooms={}

dl=300  #下次轮询间隔的秒数
for i in range(120):
    try:
        new_func()
    except EOFError:
        checkedRooms={}
    nx=datetime.now()+timedelta(seconds=dl)
    print(f'第 {i+1} 次结束,下一次时间：'+nx.strftime('%H:%M:%S'))
    time.sleep(dl)
